#include "DeadCodeEli.h"

void DeadCodeEli::execute()
{
    // printf("1\n");
    for (auto &func : this->module->get_functions())
    {
        if (func->get_basic_blocks().empty())
        {
            continue;
        }
        else
        {
            // printf("2\n");
            func_ = func;
            // printf("3\n");
            RDominateTree rdt(module);
            rdt.execute();
            mark();
            // printf("4\n");
            sweep();
            // printf("5\n");
        }
    }

    return;
}

void DeadCodeEli::mark()
{
    std::queue<Instruction *> worklist;
    for (auto bb : func_->get_basic_blocks())
    {
        for (auto curinst : bb->get_instructions())
        {
            if (curinst->is_ret() || curinst->is_call() || curinst->is_store())
            {
                worklist.push(curinst);
                mark_of_ins[curinst] = true;
            }
            else
            {
                mark_of_ins[curinst] = false;
            }
        }
    }
    while (!worklist.empty())
    {
        auto curinst = worklist.front();
        worklist.pop();
        for (auto curop : curinst->get_operands())
        {
            auto def = dynamic_cast<Instruction *>(curop);
            if (def != nullptr && mark_of_ins[def] == false)
            {
                worklist.push(def);
                mark_of_ins[def] = true;
            }
            auto bb = dynamic_cast<BasicBlock *>(curop);
            if (bb != nullptr)
            {
                auto j = bb->get_terminator();
                if (mark_of_ins[j] == false)
                {
                    mark_of_ins[j] = true;
                    worklist.push(j);
                }
            }
        }
        for (auto bb : curinst->get_parent()->get_rdom_frontier())
        {
            auto j = bb->get_terminator();
            if (mark_of_ins[j] == false)
            {
                mark_of_ins[j] = true;
                worklist.push(j);
            }
        }
    }
    simplify_func();
}

void DeadCodeEli::sweep()
{
    std::vector<Instruction *> to_delete;
    for (auto bb : func_->get_basic_blocks())
    {
        to_delete.clear();
        for (auto curinst : bb->get_instructions())
        {
            if (mark_of_ins[curinst] == false)
            {
                if (curinst->is_br() && curinst->get_num_operand() == 3)
                {
                    // printf("6\n");
                    auto jump = get_nearest_marked_postdominator(curinst->get_parent());
                    // printf("7\n");
                    curinst->remove_use_of_ops();
                    // printf("8\n");
                    curinst->set_operand(0, jump);
                    // printf("9\n");
                }
                if (!(curinst->is_br() && curinst->get_num_operand() == 1))
                {
                    to_delete.push_back(curinst);
                }
            }
        }
        for (auto del : to_delete)
        {
            bb->delete_instr(del);
        }
    }
}

BasicBlock *DeadCodeEli::get_nearest_marked_postdominator(BasicBlock *bb)
{
    std::queue<BasicBlock *> q;
    std::map<BasicBlock *, bool> vis;
    auto rdoms = bb->get_rdoms();
    for (auto curbb : rdoms)
    {
        vis[curbb] = false;
    }
    q.push(bb);
    while (!q.empty())
    {
        auto curbb = q.front();
        q.pop();
        auto children = curbb->get_succ_basic_blocks();
        for (auto child : children)
        {
            if (child != bb && rdoms.find(child) != rdoms.end()&&mark_of_ins[child->get_terminator()] == true)
            {
                    return child;
            }
            if (vis[child] == false)
            {
                q.push(child);
                vis[child] = true;
            }
        }
    }
    return nullptr;
}


void DeadCodeEli::simplify_func()
{
    std::vector<BasicBlock *> to_delete;
    for (auto bb : func_->get_basic_blocks())
    {
        if (bb == func_->get_entry_block()){
            continue;
        }
        bool flag = true;
        for (auto curinst : bb->get_instructions())
        {
            if (mark_of_ins[curinst] == true)
            {
                flag = false;
                break;
            }
        }
        if (flag)
        {
            to_delete.push_back(bb);
        }
    }
    for (auto del : to_delete)
    {
        auto inst_list = del->get_instructions();
        for (auto curinst : inst_list)
        {
            del->delete_instr(curinst);
        }
        func_->remove(del);
    }
    RDominateTree nrdt(module);
    nrdt.execute();
}